#!/usr/bin/python3
# -*- coding: utf-8 -*-

"""
# @version: v1.0
# @author : wlis
# @Email : 19688513@qq.com
# @Project : g-carbon-bio
# @File : stock_data_manager.py
# @Software: PyCharm
# @time: 2025/2/4 18:53
# @description : 股票列表-A股 接口: stock_info_a_code_name 个股信息查询 接口: stock_individual_info_em
"""

import os
import time
import logging
import sqlite3
import pandas as pd
import akshare as ak
from stock_data.models.stock_data import StockData
from path_helper import PathManager  # 引入路径管理器

# 配置日志，设置日志级别和格式
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


class StockDataManager:
    BATCH_SIZE = 100  # 设置批量插入的大小为100

    def __init__(self, db_name='stock_esg_data.db'):
        # 使用 PathManager 获取数据库存储路径
        pm = PathManager()
        self.db_name = pm.get_database_path(db_name)  # 获取数据库的完整路径
        self.max_attempts = 3  # 设置最大重试次数
        self.stock_data_list = []  # 存储股票数据的列表
        self.failed_stocks = []  # 存储获取失败的股票列表
        self.create_database()  # 创建数据库及表

    def create_database(self):
        """ 创建 SQLite 数据库和股票数据表 """
        # 检查数据库是否已存在，若不存在则创建
        if not os.path.exists(self.db_name):
            logging.info(f"数据库 {self.db_name} 不存在，正在创建。")

            # 创建数据库连接并创建表
            with sqlite3.connect(self.db_name) as conn:
                cursor = conn.cursor()
                cursor.execute('''
                    CREATE TABLE IF NOT EXISTS stocks_data (
                        stock_code TEXT PRIMARY KEY,              -- 股票代码，主键
                        stock_name TEXT,                          -- 股票名称
                        total_market_value REAL,                  -- 总市值
                        circulating_market_value REAL,             -- 流通市值
                        industry TEXT,                            -- 所属行业
                        listing_date TEXT,                        -- 上市日期
                        total_shares INTEGER,                     -- 总股本
                        circulating_shares INTEGER,                -- 流通股本
                        rating_agency TEXT,                       -- 评级机构
                        rating REAL,                              -- 评级
                        rating_quarter TEXT,                      -- 评级季度
                        market_type TEXT                          -- 市场类型
                    )
                ''')
                conn.commit()  # 提交更改到数据库
                logging.info(f"数据库 {self.db_name} 和表 stocks_data 创建成功。")
        else:
            logging.info(f"数据库 {self.db_name} 已存在，无需创建。")

    def fetch_stock_data(self):
        """ 获取股票数据并保存到 SQLite 数据库 """
        logging.info("开始获取股票数据...")
        stock_info_a_code_name_df = ak.stock_info_a_code_name()  # 获取股票代码和名称
        stock_esg_zd_sina_df = ak.stock_esg_zd_sina()  # 获取 ESG 数据
        stock_esg_zd_sina_df['股票代码'] = stock_esg_zd_sina_df['股票代码'].str.replace('.SH', '').str.replace('.SZ', '')

        # 遍历股票信息，处理每只股票
        for _, row in stock_info_a_code_name_df.iterrows():
            code = row['code'].replace('.SH', '').replace('.SZ', '')  # 提取股票代码
            name = row['name']  # 股票名称
            logging.info(f"正在处理股票: {code} - {name}")

            # 获取该股票对应的ESG数据
            esg_row = stock_esg_zd_sina_df[stock_esg_zd_sina_df['股票代码'].str.contains(code, na=False)]
            rating_agency = '秩鼎' if not esg_row.empty else '-'  # 评级机构
            rating = esg_row.iloc[0]['ESG评分'] if not esg_row.empty else '-'  # ESG评分
            rating_quarter = esg_row.iloc[0]['评分日期'] if not esg_row.empty else '-'  # 评级日期

            # 获取个股详细信息
            stock_data = self.get_individual_stock_info(code, name, rating_agency, rating, rating_quarter, self.get_stock_exchange(row['code']))
            if stock_data:
                self.stock_data_list.append(stock_data)  # 添加到股票数据列表

                # 检查是否达到批量插入的大小
                if len(self.stock_data_list) >= self.BATCH_SIZE:
                    self.batch_save_stock_data_to_db()

        # 批量保存剩余的数据
        if self.stock_data_list:
            self.batch_save_stock_data_to_db()

        logging.info("股票数据获取完成。")

    @staticmethod
    def get_stock_exchange(code):
        """
        根据股票代码判断所属交易所
        :param code: 股票代码（支持字符串或整数）
        :return: 'sh'（沪市）、'sz'（深市）、'bj'（京市）或 'unknown'
        """
        code_str = str(code).zfill(6)  # 将股票代码转换为6位字符串，不足部分用零补齐

        # 沪市判断
        if code_str.startswith('60') or code_str.startswith('688'):
            return 'sh'

        # 深市判断
        elif code_str[:3] in ['000', '001', '002', '003'] or code_str.startswith('300'):
            return 'sz'

        # 京市判断（北交所）
        elif code_str[:2] in ['43', '83', '87', '88']:
            return 'bj'

        # 如果不符合以上条件，返回未知
        else:
            return 'unknown'

    def get_individual_stock_info(self, code, name, rating_agency, rating, rating_quarter, market_type, wait_time=1):
        """ 获取个股的详细信息并返回 StockData 对象 """
        attempt = 0
        while attempt < self.max_attempts:
            try:
                time.sleep(wait_time)  # 等待一段时间以避免频繁请求
                stock_info_df = ak.stock_individual_info_em(symbol=code)  # 获取个股信息

                if stock_info_df.empty:
                    logging.error(f"获取个股信息失败，返回数据为空: 股票代码 {code}")
                    raise ValueError(f"获取个股信息失败，返回数据为空: 股票代码 {code}")

                # 将股票信息转换为字典格式
                stock_info_dict = stock_info_df.set_index('item')['value'].to_dict()
                market_type = market_type if market_type is not None else '-'  # 或者使用其他默认值

                # 创建 StockData 对象
                stock_data = StockData(
                    stock_code=stock_info_dict.get('股票代码'),
                    stock_name=stock_info_dict.get('股票简称'),
                    total_market_value=stock_info_dict.get('总市值'),
                    circulating_market_value=stock_info_dict.get('流通市值'),
                    industry=stock_info_dict.get('行业'),
                    listing_date=str(stock_info_dict.get('上市时间')),  # 确保这里强制转换为字符串
                    total_shares=stock_info_dict.get('总股本'),
                    circulating_shares=stock_info_dict.get('流通股'),
                    rating_agency=rating_agency,
                    rating=rating,
                    rating_quarter=rating_quarter,
                    market_type=market_type,  # 确保这个字段不会为 None
                )

                logging.info(f"成功获取股票信息: {code} - {name}")
                return stock_data
            except Exception as e:
                logging.error(f"获取股票 {code} 信息失败: {str(e)} (尝试次数: {attempt + 1})")
                self.failed_stocks.append({
                    'code': code,
                    'name': name,
                    'rating_agency': rating_agency,
                    'rating': rating,
                    'rating_quarter': rating_quarter,
                    'error': str(e)
                })
                attempt += 1
                time.sleep(5)  # 等待一段时间再重试

        logging.warning(f"最大尝试次数 {self.max_attempts} 达到，跳过股票 {code}")
        return None  # 返回 None 以指示获取失败

    def batch_save_stock_data_to_db(self):
        """ 批量将股票数据保存到 SQLite 数据库 """
        if not self.stock_data_list:
            logging.warning("没有股票数据需要保存。")
            return

        with sqlite3.connect(self.db_name) as conn:
            cursor = conn.cursor()
            for stock_data in self.stock_data_list:
                cursor.execute('''
                    INSERT INTO stocks_data (stock_code, stock_name, total_market_value,
                                            circulating_market_value, industry, listing_date,
                                            total_shares, circulating_shares, rating_agency,
                                            rating, rating_quarter, market_type)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                    ON CONFLICT(stock_code) DO UPDATE SET
                        stock_name = excluded.stock_name,
                        total_market_value = excluded.total_market_value,
                        circulating_market_value = excluded.circulating_market_value,
                        industry = excluded.industry,
                        listing_date = excluded.listing_date,
                        total_shares = excluded.total_shares,
                        circulating_shares = excluded.circulating_shares,
                        rating_agency = excluded.rating_agency,
                        rating = excluded.rating,
                        rating_quarter = excluded.rating_quarter,
                        market_type = excluded.market_type
                ''', (stock_data.stock_code,
                      stock_data.stock_name,
                      stock_data.total_market_value,
                      stock_data.circulating_market_value,
                      stock_data.industry,
                      stock_data.listing_date,
                      stock_data.total_shares,
                      stock_data.circulating_shares,
                      stock_data.rating_agency,
                      stock_data.rating,
                      stock_data.rating_quarter,
                      stock_data.market_type))  # 确保这里有12个值

            conn.commit()  # 提交事务
            logging.info(f"成功批量保存 {len(self.stock_data_list)} 条股票数据到数据库。")
            self.stock_data_list.clear()  # 清空列表以便下次使用

    def delete_database(self) -> None:
        """
        删除数据库文件
        """
        try:
            if os.path.exists(self.db_name):
                os.remove(self.db_name)
                logging.info(f"数据库 {self.db_name} 已成功删除")
            else:
                logging.warning(f"数据库 {self.db_name} 不存在，无法删除")
        except Exception as e:
            logging.error(f"删除数据库失败: {e}")

    def read_stock_data(self):
        """ 从 SQLite 数据库读取股票数据 """
        with sqlite3.connect(self.db_name) as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT * FROM stocks_data")  # 查询所有股票数据
            rows = cursor.fetchall()  # 获取查询结果

        stock_data_list = []
        for row in rows:
            # 将查询结果转换为 StockData 对象
            stock_data = StockData(
                stock_code=row[0],
                stock_name=row[1],
                total_market_value=row[2],
                circulating_market_value=row[3],
                industry=row[4],
                listing_date=row[5],
                total_shares=row[6],
                circulating_shares=row[7],
                rating_agency=row[8],
                rating=row[9],
                rating_quarter=row[10],
                market_type=row[11],
            )
            stock_data_list.append(stock_data)

        logging.info("从数据库加载数据成功。")
        return stock_data_list


# 使用示例
if __name__ == "__main__":
    stock_data_manager = StockDataManager()  # 实例化 StockDataManager

    logging.info("开始获取股票数据...")
    stock_data_manager.fetch_stock_data()  # 获取股票数据

    logging.info("从数据库加载数据...")
    stock_data_list = stock_data_manager.read_stock_data()  # 读取数据库中的股票数据

    # 循环输出每个 StockData 对象的数据
    if stock_data_list:
        for stock_data in stock_data_list:
            logging.info(f"股票代码: {stock_data.stock_code}, 股票名称: {stock_data.stock_name}, "
                         f"总市值: {stock_data.total_market_value}, 流通市值: {stock_data.circulating_market_value}, "
                         f"行业: {stock_data.industry}, 上市时间: {stock_data.listing_date}, "
                         f"总股本: {stock_data.total_shares}, 流通股: {stock_data.circulating_shares}, "
                         f"评级机构: {stock_data.rating_agency}, 评级: {stock_data.rating}, "
                         f"评级季度: {stock_data.rating_quarter}, 市场类型: {stock_data.market_type}")
    else:
        logging.warning("未加载任何股票数据。")
